import math

import mindspore.nn as nn
from mindspore.nn.metrics import Accuracy
from mindspore import context
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
from mindspore.train.model import Model
from mindspore.train.serialization import load_checkpoint, load_param_into_net

from src.config import cfg
from src.dataset import MovieReview
from src.dpcnn import DPCNN


if __name__ == '__main__':

    instance = MovieReview(root_dir=cfg.data_path, maxlen=cfg.word_len, split=0.9)
    dataset = instance.create_train_dataset(batch_size=cfg.batch_size,epoch_size=cfg.epoch_size)
    batch_num = dataset.get_dataset_size()

    learning_rate = []
    warm_up = [1e-3 / math.floor(cfg.epoch_size / 5) * (i + 1) for _ in range(batch_num) for i in
               range(math.floor(cfg.epoch_size / 5))]

    shrink = [1e-3 / (16 * (i + 1)) for _ in range(batch_num) for i in range(math.floor(cfg.epoch_size * 3 / 5))]

    normal_run = [1e-3 for _ in range(batch_num) for i in
                  range(cfg.epoch_size - math.floor(cfg.epoch_size / 5) - math.floor(cfg.epoch_size * 2 / 5))]
    learning_rate = learning_rate + warm_up + normal_run + shrink

    net=DPCNN()
    param_dict = load_checkpoint(cfg.checkpoint_path)
    load_param_into_net(net, param_dict)

    opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate=learning_rate, weight_decay=cfg.weight_decay)
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)

    model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc': Accuracy()})

    config_ck = CheckpointConfig(save_checkpoint_steps=int(cfg.epoch_size*batch_num/2), keep_checkpoint_max=cfg.keep_checkpoint_max)
    time_cb = TimeMonitor(data_size=batch_num)
    ckpt_save_dir = "./ckpt"
    ckpoint_cb = ModelCheckpoint(prefix="train_dpcnn", directory=ckpt_save_dir, config=config_ck)
    loss_cb = LossMonitor()
    model.train(cfg.epoch_size, dataset, callbacks=[time_cb, ckpoint_cb, loss_cb])
    print("train success")

